import os
import glob
import warnings
import numpy as np
import mne
import scipy.io as sio
import matplotlib.pyplot as plt
import matplotlib


mne.set_log_level('ERROR')

# 解决中文显示与负号显示问题（Windows 默认可用 SimHei/微软雅黑）
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS', 'DejaVu Sans']
matplotlib.rcParams['axes.unicode_minus'] = False


def read_concat_raws(gdf_paths):
    if len(gdf_paths) == 0:
        raise FileNotFoundError("未找到任何.gdf文件用于读取")
    raw_list = []
    for p in gdf_paths:
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message=r"Highpass cutoff frequency .* greater than lowpass cutoff frequency .*",
                category=RuntimeWarning,
            )
            r = mne.io.read_raw_gdf(p, preload=True, verbose=False)
        hp = r.info.get('highpass', None)
        lp = r.info.get('lowpass', None)
        if hp is not None and lp is not None and hp > lp:
            r.info['highpass'], r.info['lowpass'] = lp, hp
        raw_list.append(r)
    return mne.io.concatenate_raws(raw_list)


def load_mat_labels(mat_path):
    md = sio.loadmat(mat_path, squeeze_me=True, struct_as_record=False)
    for key in ['labels', 'y', 'Y', 'classlabel', 'label']:
        if key in md:
            lab = np.asarray(md[key]).ravel()
            break
    else:
        lab = None
        for v in md.values():
            if isinstance(v, np.ndarray) and v.ndim == 1 and v.dtype.kind in ('i', 'u'):
                uniq = np.unique(v)
                if 1 <= uniq.size <= 3:
                    lab = v
                    break
        if lab is None:
            raise ValueError(f"无法在{os.path.basename(mat_path)}中推断标签向量")
    uniq = np.unique(lab)
    if set(uniq.tolist()) <= {0, 1}:
        y = lab.astype(int)
    elif set(uniq.tolist()) <= {1, 2}:
        y = (lab.astype(int) - 1)
    elif set(uniq.tolist()) <= {769, 770}:
        y = (lab.astype(int) == 770).astype(int)
    else:
        raise ValueError(f"不支持的标签取值：{uniq}")
    return y


def build_events_from_onsets_and_labels(raw, labels, prefer_codes=("768", "781")):
    events_all, event_id_all = mne.events_from_annotations(raw, verbose=False)
    sel_code = None
    for code in prefer_codes:
        if code in event_id_all:
            sel_code = code
            break
    if sel_code is None:
        ev = events_all[:len(labels)]
    else:
        ev = events_all[events_all[:, 2] == event_id_all[sel_code]]
        if ev.shape[0] < len(labels):
            ev = events_all[:len(labels)]
    if ev.shape[0] < len(labels):
        raise ValueError(f"试次起点数量({ev.shape[0]})少于标签数量({len(labels)})，无法对齐")
    ev = ev[:len(labels)].copy()
    ev[:, 2] = np.where(labels == 0, 769, 770)
    event_id_map = {'left': 769, 'right': 770}
    return ev, event_id_map


def pick_c3_cz_c4_c5_c6(raw):
    ch_names = raw.info['ch_names']
    lower_to_orig = {c.lower(): c for c in ch_names}
    # 需要的通道，固定顺序：C3, C5, Cz, C4, C6
    targets = ['c3', 'c5', 'cz', 'c4', 'c6']
    resolved = []
    for t in targets:
        resolved.append(lower_to_orig.get(t))
    for i, t in enumerate(targets):
        if resolved[i] is not None:
            continue
        cand = [c for c in ch_names if t in c.lower()]
        if len(cand) > 0:
            resolved[i] = cand[0]
    if any(r is None for r in resolved):
        missing = [targets[i].upper() for i, r in enumerate(resolved) if r is None]
        raise ValueError(f"未找到所需通道: {missing}；可用通道: {ch_names}")
    raw.pick_channels(resolved, ordered=True, verbose=False)


def compute_mean_power_curves(epochs, sfreq):
    # epochs: (n_trials, 5, n_times) after filtering 8-30 Hz, order: C3,C5,Cz,C4,C6
    X = epochs.get_data()  # already filtered
    power = X ** 2  # 瞬时功率
    mean_power = power.mean(axis=0)  # shape: (5, n_times)
    times = epochs.times  # seconds
    # 左侧平均：C3与C5；右侧平均：C4与C6
    left_avg = (mean_power[0] + mean_power[1]) / 2.0
    right_avg = (mean_power[3] + mean_power[4]) / 2.0
    return times, left_avg, right_avg


def main():
    base_dir = os.path.join(os.path.dirname(__file__), '2b')
    dir_train = os.path.join(base_dir, 'train')
    dir_train_label = os.path.join(base_dir, 'train_label')

    train_gdf = sorted(glob.glob(os.path.join(dir_train, '*.gdf')))
    train_mat = {os.path.splitext(os.path.basename(p))[0]: p for p in sorted(glob.glob(os.path.join(dir_train_label, '*.mat')))}
    if len(train_gdf) == 0:
        raise FileNotFoundError("未找到训练集.gdf文件")

    # 受试者分组（如 B01、B02... 由文件名前三位确定）
    subj_to_files = {}
    for gdf_path in train_gdf:
        stem = os.path.splitext(os.path.basename(gdf_path))[0]
        subj = stem[:3]
        subj_to_files.setdefault(subj, []).append(gdf_path)

    results = []
    for subj, files in sorted(subj_to_files.items()):
        epochs_all = None
        for gdf_path in sorted(files):
            r = read_concat_raws([gdf_path])
            pick_c3_cz_c4_c5_c6(r)
            r.set_eeg_reference('average', projection=True, verbose=False)
            r.filter(8, 30, fir_design='firwin', verbose=False)
            stem = os.path.splitext(os.path.basename(gdf_path))[0]
            mat_path = train_mat.get(stem)
            if mat_path is None:
                raise FileNotFoundError(f"未找到与{stem}匹配的标签.mat")
            y = load_mat_labels(mat_path)
            events, _ = build_events_from_onsets_and_labels(r, y)
            sel = events[:, 2] == 770  # 右手
            if sel.sum() == 0:
                continue
            epochs = mne.Epochs(r, events[sel], event_id={'right': 770}, tmin=0.0, tmax=9.0,
                                baseline=(0, 0), preload=True, verbose=False)
            epochs_all = epochs if epochs_all is None else mne.concatenate_epochs([epochs_all, epochs])

        if epochs_all is None:
            print(f"[警告] 受试者{subj}没有右手试次，跳过。")
            continue

        times, c3_curve, c4_curve = compute_mean_power_curves(epochs_all, epochs_all.info['sfreq'])
        diff = np.abs(c3_curve - c4_curve)
        idx_max = int(np.argmax(diff))
        t_max = float(times[idx_max])
        sfreq = epochs_all.info['sfreq']
        win = int(round(sfreq * 1.0))
        win = max(win, 1)
        best_mean = -np.inf
        best_start = 0
        for start in range(0, diff.size - win + 1):
            m = diff[start:start+win].mean()
            if m > best_mean:
                best_mean = m
                best_start = start
        t_start = float(times[best_start])
        t_end = float(times[min(best_start + win - 1, diff.size - 1)])

        results.append((subj, t_max, t_start, t_end))

        # 绘制并保存每位受试者的图
        plt.figure(figsize=(8, 3))
        plt.plot(times, c3_curve, label='左侧(C3+C5)/2', linestyle='--')
        plt.plot(times, c4_curve, label='右侧(C4+C6)/2', linestyle='-')
        plt.axvline(t_max, color='r', linestyle=':', label=f'|C3-C4|最大@{t_max:.2f}s')
        plt.axvspan(t_start, t_end, color='orange', alpha=0.2, label=f'最佳1秒区间[{t_start:.2f},{t_end:.2f}]s')
        plt.xlabel('t / s')
        plt.ylabel('平均功率 (a.u.)')
        plt.title(f'右手想象期间 C3 与 C4 的平均功率（{subj}）')
        plt.legend(loc='best')
        out_path = os.path.join(os.path.dirname(__file__), f'c3_c4_right_power_{subj}.png')
        plt.tight_layout()
        plt.savefig(out_path, dpi=150)
        plt.close()
        print(f"[{subj}] 绘图已保存：{out_path}")

    # 终端输出每位受试者的最佳时间段
    print("\n=== 受试者依赖的最佳时间段（右手） ===")
    for subj, t_max, t_start, t_end in results:
        print(f"{subj}: |C3-C4|最大差异@{t_max:.3f}s, 最佳1秒区间[{t_start:.3f}, {t_end:.3f}]s")

    times, c3_curve, c4_curve = compute_mean_power_curves(epochs_all, epochs_all.info['sfreq'])
    diff = np.abs(c3_curve - c4_curve)
    # 最大瞬时差异时刻
    idx_max = int(np.argmax(diff))
    t_max = float(times[idx_max])
    # 最佳1秒区间（滑窗平均差异最大）
    sfreq = epochs_all.info['sfreq']
    win = int(round(sfreq * 1.0))
    if win < 1:
        win = 1
    # 有效滑动范围
    best_mean = -np.inf
    best_start = 0
    for start in range(0, diff.size - win + 1):
        m = diff[start:start+win].mean()
        if m > best_mean:
            best_mean = m
            best_start = start
    t_start = float(times[best_start])
    t_end = float(times[min(best_start + win - 1, diff.size - 1)])

    # 绘图
    plt.figure(figsize=(8, 3))
    plt.plot(times, c3_curve, label='C3', linestyle='--')
    plt.plot(times, c4_curve, label='C4', linestyle='-')
    plt.axvline(t_max, color='r', linestyle=':', label=f'|C3-C4|最大@{t_max:.2f}s')
    plt.axvspan(t_start, t_end, color='orange', alpha=0.2, label=f'最佳1秒区间[{t_start:.2f},{t_end:.2f}]s')
    plt.xlabel('t / s')
    plt.ylabel('平均功率 (a.u.)')
    plt.title('右手想象期间 C3 与 C4 的平均功率')
    plt.legend(loc='best')
    out_path = os.path.join(os.path.dirname(__file__), 'c3_c4_right_power.png')
    plt.tight_layout()
    plt.savefig(out_path, dpi=150)
    plt.close()

    # 上述循环中已输出各受试者结果


if __name__ == '__main__':
    main()


